import sys, os, time
from PyQt4 import QtGui, QtCore
try:
    import VISA_Agilent_8753c_NA as VV
except:
    pass
import Coil_sp_to_RLM as CC

IdRole = QtCore.Qt.UserRole

class RenderArea(QtGui.QWidget):
    def __init__(self, parent = None):
        QtGui.QWidget.__init__(self, parent)
        self.setFocusPolicy(QtCore.Qt.ClickFocus)
        self.pixmap = QtGui.QPixmap()
        self.pixname = ""
        self.setBackgroundRole(QtGui.QPalette.Base)
        self.image = QtGui.QImage()
        self.scribbling = ""
        self.lastPoint = None

    def resizeEvent(self, event):
        if self.width() > self.image.width() or self.height() > self.image.height():
            newWidth = max(self.width() + 128, self.image.width())
            newHeight = max(self.height() + 128, self.image.height())
            self.resizeImage(self.image, QtCore.QSize(newWidth, newHeight))
            self.update()
        QtGui.QWidget.resizeEvent(self, event)
        
    def resizeImage(self, image, newSize):
        if image.size() == newSize:
            return
        newImage = QtGui.QImage(newSize, QtGui.QImage.Format_ARGB32)
        newImage.fill(QtGui.qRgba(255, 255, 255,0))
        painter = QtGui.QPainter()
        painter.begin(newImage)
        painter.drawImage(QtCore.QPoint(0, 0), image)
        painter.end()
        self.image = newImage
    def clearImage(self):
        self.image.fill(QtGui.qRgba(255, 255, 255,0))
        self.modified = True
        self.update()        
    def minimumSizeHint(self):
        return QtCore.QSize(100, 100)

    def sizeHint(self):
        return QtCore.QSize(850, 600)

    def setPixmap(self, pixname):
        if self.pixmap.load(pixname):
            self.pixname = pixname
            print "%s loaded"%pixname
        else:
            if not os.path.exists(pixname):
                print "%s picture not found"%pixname
            else:
                print "%s exists, but can't be loaded"%pixname
        self.update()
    
    def paintEvent(self, event):
        painter = QtGui.QPainter()
        painter.begin(self)
        painter.setRenderHints(QtGui.QPainter.Antialiasing | QtGui.QPainter.SmoothPixmapTransform)
        painter.drawPixmap(0, 0, self.pixmap)
        painter.drawImage(0,0,self.image)
        painter.end()

    def keyPressEvent(self, event):
        key = event.key()
        if key == QtCore.Qt.Key_X:
            self.scribbling = "X"
            self.lastPoint = self.mapFromGlobal(QtGui.QCursor.pos())
            self.setMouseTracking(True)
            self.clearImage()
            self.update()
            print "zooming x axis"
        elif key == QtCore.Qt.Key_Y:
            self.scribbling = "Y"
            self.lastPoint = self.mapFromGlobal(QtGui.QCursor.pos())
            self.setMouseTracking(True)
            self.clearImage()
            self.update()
            print "zooming y axis"
        elif key == QtCore.Qt.Key_Escape:
            self.scribbling = ""
            self.setMouseTracking(False)
            self.clearImage()
            self.update()
            print "zooming cancelled"
        elif key == QtCore.Qt.Key_F:
            self.scribbling = ""
            self.setMouseTracking(False)
            self.clearImage()
            painter = QtGui.QPainter()
            testimg = QtGui.QImage.fromData(CC.drawData(CC.array(range(1,100)),CC.array(range(1,100))).makeChart2(CC.PNG))
            painter.begin(self.image)
            painter.drawImage(0,0,testimg)
            painter.end()
            self.update()
            print "zoom to best fit"
    def mousePressEvent(self, event):
        if event.button() == QtCore.Qt.LeftButton and self.scribbling == "":
            self.lastPoint = self.mapFromGlobal(QtGui.QCursor.pos())
            self.scribbling = "XY"
            self.setMouseTracking(True)
            print "zoom X and y axis"
    def mouseMoveEvent(self, event):
        if (event.buttons() & QtCore.Qt.LeftButton) or self.scribbling in ["X", "Y", "XY"]:
            self.clearImage()
            painter = QtGui.QPainter()
            painter.begin(self.image)
            if self.scribbling == "X":
                painter.drawLine(self.lastPoint.x(),0, self.lastPoint.x(),self.image.height())
                painter.drawLine(event.pos().x(),0, event.pos().x(),self.image.height())
                painter.drawLine(self.lastPoint, event.pos())
            elif self.scribbling == "Y":
                painter.drawLine(0,self.lastPoint.y(),self.image.width(), self.lastPoint.y())
                painter.drawLine(0,event.pos().y(),self.image.width(), event.pos().y())
                painter.drawLine(self.lastPoint, event.pos())
            elif self.scribbling == "XY":
                painter.drawLine(self.lastPoint.x(),self.lastPoint.y(),self.lastPoint.x(), event.pos().y())
                painter.drawLine(self.lastPoint.x(), event.pos().y(),event.pos().x(), event.pos().y())
                painter.drawLine(event.pos().x(), event.pos().y(),event.pos().x(),self.lastPoint.y())
                painter.drawLine(event.pos().x(),self.lastPoint.y(),self.lastPoint.x(),self.lastPoint.y())                
            painter.end()
    def mouseReleaseEvent(self, event):
        """(50,35) left top corner, (750,455) right bottom corner"""
        if event.button() == QtCore.Qt.LeftButton and self.scribbling:
            print "Left Released, mode: %s, from:%s, to:%s"%(self.scribbling, self.lastPoint, event.pos())
            self.scribbling = ""
        self.clearImage()
        self.setMouseTracking(False)

class myLineBox(QtGui.QLineEdit):
    def __init__(self, parent = None):
        QtGui.QLineEdit.__init__(self, parent)
        self.parentWin = parent
    def keyPressEvent(self, event):
        key = event.key()

        if key in [QtCore.Qt.Key_Down, QtCore.Qt.Key_Up]:
            filelist = os.listdir(self.parentWin.workingfolder)
            if key == QtCore.Qt.Key_Up:
                filelist = filelist[::-1]
            nowvalue = str(self.text())
            for i in range(len(filelist)):
                if not filelist[i].lower().endswith(".s2p"):
                    continue
                if nowvalue in filelist[i]:
                    bFound = False
                    for j in range(i+1, len(filelist)):
                        if filelist[j].lower().endswith(".s2p"):
                            self.setText(filelist[j])
                            self.selectAll()
                            bFound = True
                            break
                    if bFound: break
        QtGui.QLineEdit.keyPressEvent(self, event)

class Window(QtGui.QWidget):
    def __init__(self, parent = None):
        QtGui.QWidget.__init__(self, parent)
        self.setWindowTitle(self.tr("Basic Drawing"))

        self.s2pname = ""
        self.renderArea = RenderArea()
        self.workingfolder = "."
        self.twoPorts = False

        self.folderComboBox = self.createComboBox(QtCore.QDir.currentPath())
        folderLabel = QtGui.QLabel(self.tr("Working Folder:"))
        setfolderButton = self.createButton(self.tr("&Browse..."), self.setFolder)
        folderLabel.setBuddy(setfolderButton)

        self.pixLineEdit = myLineBox(self)
        self.pixLabel = QtGui.QLabel(self.tr("Plot &Name:"))
        self.pixLabel.setBuddy(self.pixLineEdit)
        
        plotoptions = (("R",1),("L",2),("Q",3),("K",4),)
        self.plotComboBox = QtGui.QComboBox()
        for plotname,numid in plotoptions:
            self.plotComboBox.addItem(self.tr(plotname), QtCore.QVariant(numid))
        self.plotLabel = QtGui.QLabel(self.tr("&Plot types:"))
        self.plotLabel.setBuddy(self.plotComboBox)

        self.portsCheckBox = QtGui.QCheckBox(self.tr("Two Po&rts:"))

        genS2pButton = QtGui.QPushButton(self.tr("Generate .s2p"))
        resetButton = QtGui.QPushButton(self.tr("reset NA"))
        writeCaliButton = QtGui.QPushButton(self.tr("write Calib to NA"))
        loadCaliButton = QtGui.QPushButton(self.tr("load Calib from NA"))
        clearButton = QtGui.QPushButton(self.tr("clear log files"))

        self.connect(self.portsCheckBox, QtCore.SIGNAL("toggled(bool)"),
                     self.portsChanged)
        self.connect(self.pixLineEdit, QtCore.SIGNAL("textChanged(const QString &)"),
                     self.s2pNameChanged)
        self.connect(self.folderComboBox, QtCore.SIGNAL("textChanged(const QString &)"),
                     self.folderChanged)
        self.connect(self.plotComboBox, QtCore.SIGNAL("activated(int)"),
                     self.plotChanged)

        self.connect(genS2pButton, QtCore.SIGNAL("clicked()"),
                     self.genS2p)
        self.connect(resetButton, QtCore.SIGNAL("clicked()"),
                     self.resetNA)
        self.connect(writeCaliButton, QtCore.SIGNAL("clicked()"),
                     self.writeCali2NA)
        self.connect(loadCaliButton, QtCore.SIGNAL("clicked()"),
                     self.loadCaliFromNA)
        self.connect(clearButton, QtCore.SIGNAL("clicked()"),
                     self.clearLogFiles)
        
        mainLayout = QtGui.QGridLayout()
        mainLayout.addWidget(self.renderArea, 0, 0, 1, 4)
        mainLayout.addWidget(folderLabel, 1,0)
        mainLayout.addWidget(self.folderComboBox,1,1,1,2)
        mainLayout.addWidget(setfolderButton,1,3)
        mainLayout.addWidget(self.pixLabel,2,0)
        mainLayout.addWidget(self.pixLineEdit, 2,1,1,3)
        mainLayout.addWidget(self.plotLabel, 3,0)
        mainLayout.addWidget(self.plotComboBox,3,1,1,3)
        mainLayout.addWidget(self.portsCheckBox, 4,0)
        mainLayout.addWidget(genS2pButton, 4,1,1,3)
        mainLayout.addWidget(resetButton, 5,0)
        mainLayout.addWidget(writeCaliButton, 5,1)
        mainLayout.addWidget(loadCaliButton, 5,2)
        mainLayout.addWidget(clearButton, 5,3)
        
        self.setLayout(mainLayout)
        self.plotChanged()
        self.portsChanged(False)
        self.renderArea.setFocus()
    def createButton(self, text, member):
        button = QtGui.QPushButton(text)
        self.connect(button, QtCore.SIGNAL("clicked()"), member)
        return button
    def createComboBox(self, text=""):
        comboBox = QtGui.QComboBox()
        comboBox.setEditable(True)
        comboBox.addItem(text)
        comboBox.setSizePolicy(QtGui.QSizePolicy.Expanding, QtGui.QSizePolicy.Preferred)
        return comboBox    
    def setFolder(self):
        directory = QtGui.QFileDialog.getExistingDirectory(self, self.tr("Find Files"),
                                                           QtCore.QDir.currentPath())
        if len(str(directory))==0:
            directory = "."
        self.folderComboBox.addItem(directory)
        self.folderComboBox.setCurrentIndex(self.folderComboBox.currentIndex() + 1)
    def clearLogFiles(self):
        reply = QtGui.QMessageBox.question(self, self.tr("delete ?"),
                                   self.tr("delete all *.txt and *.png files ?"), QtGui.QMessageBox.Yes,
                                   QtGui.QMessageBox.No)
        if reply == QtGui.QMessageBox.Yes:
            filelist = os.listdir(".")
            for f in filelist:
                f = f.lower()
                if ".s2p_" in f:
                    if f.endswith(".txt") or f.endswith(".png"):
                        os.remove(f)
                    
    def resetNA(self):
        try:
            VV.reset()
        except:
            pass
    def writeCali2NA(self):
        try:
            ca = VV.readCaliFromFile()
            VV.writeCali(ca)
        except:
            pass
    def loadCaliFromNA(self):
        try:
            ca = VV.readCali()
            if ca:
                VV.writeCaliToFile(ca)
        except:
            pass
            
    def plotChanged(self):
        value = self.plotComboBox.itemData(self.plotComboBox.currentIndex(), IdRole).toInt()[0]
        comboOptions = {1:"_RvsF.png",2:"_LvsF.png",
                        3:"_QvsF.png",4:"_KvsF.png"}
        pixname = os.path.join(self.workingfolder,self.s2pname)+comboOptions[value]
        if not os.path.exists(os.path.join(self.workingfolder,self.s2pname)):
            print "%s not found"%self.s2pname
            return
        if not os.path.exists(pixname) or os.path.getmtime(os.path.join(self.workingfolder,self.s2pname))>=os.path.getmtime(pixname):
            if os.path.isfile(os.path.join(self.workingfolder,self.s2pname)):
                print "%s is newer, re-calculating parameters..."%self.s2pname
                self.drawS2p()
            else:
                return
        self.renderArea.setPixmap(pixname)
        self.update()
        
    def portsChanged(self, bTwoPorts):
        self.twoPorts = bTwoPorts
        if os.path.exists(self.s2pname):
            self.drawS2p()
        self.plotChanged()
    def folderChanged(self, foldername):
        self.pixLineEdit.setText("")
        self.workingfolder = str(self.folderComboBox.currentText())
        print "%s is found: %s"%(self.workingfolder, os.path.isdir(self.workingfolder))
    def s2pNameChanged(self, filename):
        self.s2pname = str(filename)
        if len(filename)>2:
            filelist = os.listdir(".")
            for f in filelist:
                if not f.lower().endswith(".s2p"):
                    continue
                if f.startswith(filename):
                    self.pixLineEdit.setText(f)
                    self.pixLineEdit.setSelection(len(filename), len(f)-len(filename))
                    self.s2pname = f
                    break
        self.plotChanged()
    def genS2p(self):
            if os.path.exists(self.s2pname):
                reply = QtGui.QMessageBox.question(self, self.tr("overwrite ?"),
                                           self.tr("file %s already exists, overwrite ?"%self.s2pname), QtGui.QMessageBox.Yes,
                                           QtGui.QMessageBox.No)
                if reply == QtGui.QMessageBox.No:
                    return
            VV.genS2p(self.s2pname)
            self.plotChanged()
    def drawS2p(self):
        CC.drawFile(os.path.join(self.workingfolder,self.s2pname), self.twoPorts)
    def paintEvent(self,event):
        self.renderArea.update()

if __name__ == "__main__":
    app = QtGui.QApplication(sys.argv)
    window = Window()
    window.show()
    sys.exit(app.exec_())
